import os
import shutil
import matplotlib.pyplot as plt

from glob import glob
from utils import get_json_points, get_map_from_points, get_map_num, get_fidtm_map_num


# 测试数据、方法信息

data_name = 'RGBTCC'

gt_path = "../datasets/rgbtcc/test"
rgb_path = "../datasets/rgbtcc/test"
tir_path = "../datasets/rgbtcc/test"
gt_list = glob(os.path.join(gt_path, '*.json'))

method_results = {
    'CAGNet': f"../other_methods/MIANet/map_{data_name}",
    'MC3Net': f"../other_methods/MC3Net/map_{data_name}",
    'GETANet': f"../other_methods/GETANet/map_{data_name}",
    'DEFNet': f"../other_methods/DEFNet/map_{data_name}",
    'MSDTrans': f"../other_methods/MSDTrans/map_{data_name}",
    'FIDTM': f"/data/store1/nzd/mamba/prj/fidtm_tir_rgbtcc/map_{data_name}",
}

save_path = './1_visul_comparison_rgbtcc'
if os.path.exists(save_path):
    shutil.rmtree(save_path)
os.makedirs(save_path)


for i, gt_file in enumerate(gt_list):
    # print(f">>>>>{i}", end="\r")
    
    # 1.获取原始数据
    gt_name = os.path.basename(gt_file)
    et_name = gt_name.replace("_GT.json", "_RGB.pt")
    rgb_file = os.path.join(rgb_path, gt_name.replace('_GT.json', '_RGB.jpg'))
    tir_file = os.path.join(tir_path, gt_name.replace('_GT.json', '_T.jpg'))

    rgb_img = plt.imread(rgb_file)
    tir_img = plt.imread(tir_file)
    image_shape = rgb_img.shape[:2] # (h, w)
    
    points = get_json_points(gt_file)
    gt_num = len(points)
    gt_map = get_map_from_points(points, image_shape, sigma=8)

    # 2.保存对比结果
    # plt参数
    fixed = 3
    row = 1
    column = len(method_results) + fixed
    font_properties = {
        'fontsize': 6,
        'fontweight': 'normal',
        'fontfamily': 'serif',
        'color': 'darkred'
    }

    # 画图
    fig, axes = plt.subplots(row, column, dpi=800, tight_layout=True, frameon=False)

    axes[0].imshow(rgb_img)
    if len(points) > 0:
        axes[0].scatter([p[0] for p in points], [p[1] for p in points], s=0.05, c="red")
    axes[0].axis("off")
    axes[0].set_title(f"{et_name}", fontdict=font_properties)
    
    axes[1].imshow(tir_img)
    if len(points) > 0:
        axes[1].scatter([p[0] for p in points], [p[1] for p in points], s=0.05, c="red")
    axes[1].axis("off")
    # axes[1].set_title(f"{image_name}", fontdict=font_properties)

    axes[2].imshow(gt_map, cmap="jet")
    axes[2].axis("off")
    axes[2].set_title(f"{gt_num}", fontdict=font_properties)

    # 3.获取方法结果
    for i, method in enumerate(method_results.keys()):
        et_file = os.path.join(method_results[method], et_name)

        if not os.path.exists(et_file):
            print(et_file)
            continue
        # map, num
        if method == 'FIDTM':
            et_map, et_num = get_fidtm_map_num(et_file)
        else:
            et_map, et_num = get_map_num(et_file, image_shape)
        
        axes[i+fixed].imshow(et_map, cmap="jet")
        axes[i+fixed].axis("off")
        axes[i+fixed].set_title(f"{method}:{et_num}", fontdict=font_properties)

    # plt.show()
    plt.savefig(os.path.join(save_path, et_name.split('.')[0]), bbox_inches='tight', pad_inches=0)
    plt.close()



for i, gt_file in enumerate(gt_list):
    # print(f">>>>>{i}", end="\r")
    
    # 1.获取原始数据
    gt_name = os.path.basename(gt_file)
    et_name = gt_name.replace("_GT.json", "_RGB.pt")
    rgb_file = os.path.join(rgb_path, gt_name.replace('_GT.json', '_RGB.jpg'))
    tir_file = os.path.join(tir_path, gt_name.replace('_GT.json', '_T.jpg'))

    rgb_img = plt.imread(rgb_file)
    tir_img = plt.imread(tir_file)
    image_shape = rgb_img.shape[:2] # (h, w)
    
    points = get_json_points(gt_file)
    gt_num = len(points)
    gt_map = get_map_from_points(points, image_shape, sigma=8)

    # 2.保存对比结果
    # plt参数
    fixed = 3
    row = 1
    column = len(method_results) + fixed
    font_properties = {
        'fontsize': 6,
        'fontweight': 'normal',
        'fontfamily': 'serif',
        'color': 'darkred'
    }

    # 画图
    fig, axes = plt.subplots(row, column, dpi=800, tight_layout=True, frameon=False)

    axes[0].imshow(rgb_img)
    if len(points) > 0:
        axes[0].scatter([p[0] for p in points], [p[1] for p in points], s=0.05, c="red")
    axes[0].axis("off")
    axes[0].set_title(f"{et_name}", fontdict=font_properties)
    
    axes[1].imshow(tir_img)
    if len(points) > 0:
        axes[1].scatter([p[0] for p in points], [p[1] for p in points], s=0.05, c="red")
    axes[1].axis("off")
    # axes[1].set_title(f"{image_name}", fontdict=font_properties)

    axes[2].imshow(gt_map, cmap="jet")
    axes[2].axis("off")
    axes[2].set_title(f"{gt_num}", fontdict=font_properties)

    # 3.获取方法结果
    for i, method in enumerate(method_results.keys()):
        et_file = os.path.join(method_results[method], et_name)

        if not os.path.exists(et_file):
            print(et_file)
            continue
        # map, num
        if method == 'FIDTM':
            et_map, et_num = get_fidtm_map_num(et_file)
        else:
            et_map, et_num = get_map_num(et_file, image_shape)
        
        axes[i+fixed].imshow(et_map, cmap="jet")
        axes[i+fixed].axis("off")
        axes[i+fixed].set_title(f"{method}:{et_num}", fontdict=font_properties)

    # plt.show()
    plt.savefig(os.path.join(save_path, et_name.split('.')[0]), bbox_inches='tight', pad_inches=0)
    plt.close()
